import tensorflow as tf
import numpy as np
import matplotlib.pyplot as plt
import os
from LSTMModel.dataReader import getJsonData,DataReader
from LSTMModel.config import lstm_config
from LSTMModel.utils import getAccuracy,getConfusionMatrix,plot_confusion_matrix,softmax_2D


class Model():
    # 预测类
    def __init__(self,weight_path,label_path,graph=None):
        self._sess=None              # 会话
        self._inputs_data=None       # 输入接口
        self._outputs=None           # 输出接口
        self.label2id=None
        self.id2label=None
        self._dropout=None

        if graph is None:
            graph=tf.Graph()
        self._graph=graph

        self.loadWeight(weight_path)
        self.loadLabelId(label_path)

    def loadWeight(self,weight_path):
        # 加载权重
        self._sess=tf.compat.v1.Session(graph=self._graph)
        with self._sess.as_default():
            with self._graph.as_default():
                new_saver = tf.compat.v1.train.import_meta_graph(weight_path+".meta")    # 加载元图
                new_saver.restore(self._sess,weight_path)
                graph=tf.get_default_graph()        # 获取计算图
                #########查看图节点##########
                # for i in graph.as_graph_def().node:
                #     print(i.name)
                # exit()
                ###########################
                self._inputs_data=graph.get_tensor_by_name("inputs_data:0")   # 获取输入接口
                self._outputs=graph.get_tensor_by_name("prediction/BiasAdd:0")        # 获取输出接口
                self._dropout=graph.get_tensor_by_name("dropout:0")

    def loadLabelId(self,label_path):
        self.label2id,self.id2label=getJsonData(label_path,
                                                  ['label2id','id2label'])

    def predict(self,seq_data):
        # 预测
        with self._sess.as_default():
            with self._sess.graph.as_default():
                try:
                    prediction_prob=self._sess.run(self._outputs,feed_dict={self._inputs_data:seq_data,
                                                                            self._dropout:1.0})
                    prediction_prob=softmax_2D(prediction_prob)
                    max_prob_index=np.argmax(prediction_prob,axis=-1)
                    prediction_word=[(self.id2label[str(word_index)],prediction_prob[index,word_index]) for index,word_index in enumerate(max_prob_index)]
                except:
                    raise ValueError("data format error!")
                return prediction_word

    def close(self):
        self._sess.close()


if __name__=="__main__":
    model=Model(weight_path="/home/lisen/tool/PyProjects/唇语识别/weight_files/LSTM_weight/weight_file_v2_128_1_09/weight_ckpt",
                label_path="/home/lisen/tool/PyProjects/唇语识别/dataset/dataset_word/dataset_v2/ids_three.json")

    test_dir = os.path.join(lstm_config.data_dir, "test")  # 测试集目录
    ids_path = os.path.join(lstm_config.data_dir, "ids_three.json")  # 标签路径
    # 获取数据
    test_reader = DataReader(data_dir=test_dir,
                             ids_path=ids_path,
                             feature_len=lstm_config.feature_len,
                             seq_max_len=lstm_config.max_seq_len)
    test_data, test_labels = test_reader.getAllData()
    test_labels=np.argmax(test_labels,axis=1)

    # 正确率
    prediction_word=model.predict(test_data)
    model.close()
    prediction=[model.label2id[i[0]] for i in prediction_word]
    acc = getAccuracy(test_labels, prediction)
    cm=getConfusionMatrix(test_labels, prediction)
    print("test accurcy:{}%".format(acc * 100))
    print("confusion matrix:\n{}".format(cm))
    label_name=[test_reader.id2label[str(i)] for i in range(3)]
    plot_confusion_matrix(cm,label_name,"confusion_matrix")
    plt.show()